#!/usr/bin/python

#Programmer: Aaron Steele
#Date: 1/17/2013
#Purpose: Extract the nucleotide sequences from the PFam seed file and the
#		  web crawler output.
#Usage: pfam-seq-extract <web-cralwer-output> <Pfam.seed>
 		  
#TODO Create a direcotry for each Protein family to be stored in
#TODO Handle the input and output in a more dynamic fashion
#TODO Double check pfam_identifier naming methods. PF000#.## <- might be wrong #

import sys
import string
	
#Read the input parameters
nucfile = sys.argv[1]
pfamfile = sys.argv[2]
out_dir = sys.argv[3]

id = list()
seq = list()

#Read in all the ids and sequences from the file
pfam_convert(nucfile,id, seq)

print "Length: " + str(len(id))
#Open the pfam seed file
pfamIn = open(pfamfile, "r")
outfile = open(out_dir + "emtpyfile", "w")
cur_id = -1
i = 0;

while 1:
	line = pfamIn.readline()
	line = line.rstrip()
	#print line

	#If the last line exit the loop
	if not line:
		break

	if line.startswith("# STOCKHOLM 1.0"):
		#This is a new protein family
		outfile.close()
		i = i+1			
			
		#Get the protein family id and open new file
		line = pfamIn.readline()
		line = pfamIn.readline()
		line = line.rstrip()
		split_string = line.split()
		outfile = open(out_dir + split_string[2], "w")
				
	else:
		if line.startswith("#=GS"):
			split_string = line.split()
			if split_string[2] == "AC":
				#Get the sequence id 
				tokens = split_string[3].split(".")
				cur_id = tokens[0]
			
				#find the id in the id list
				for j in range(len(id)):
					if id[j] == cur_id:
						outfile.write(">")
						outfile.write(split_string[1])
						outfile.write("\n")
						outfile.write(seq[j])
						outfile.write("\n")
						outfile.flush()

						#Remove the id and sequence to prevent duplication and improve efficiency
						del id[j]
						del seq[j]
						break
print "Protein Families Found: " + str(i) + "\n"
					
def pfam_convert(infile, id, seq):

	print "Parsing the nuc file\n"
	#Open the input file
	input = open(infile, "r")
	
	while 1:
		line = input.readline()
		
		#If the last line exit the loop
		if not line:
			break

		line = string.replace(line, "//", "")
		line = line.rstrip()
		split_string = string.split(line,"=")
		
		if split_string[0] == "u_id":
			id.append(split_string[1])

		elif split_string[0] == "coding_seq":
			seq.append(split_string[1])

	#Close the input and output files
	print "Parsing complete\n"
	input.close();
